#!/usr/bin/python
# -*- coding: utf-8 -*-

import simplejson
import MuConnection
import time
from StringIO import StringIO
from xml.etree import ElementTree
import logging

log = logging.getLogger('MuPython.MuRestApi')

""" MuRestApi contains the basic Mu REST API commands necessary to execute 
all DoS, PVA, Studio Fuzz, and admin/maintanence functions 

Classes contained: 
    MuRestApi - This class is instantiated to run Fuzz anaysis, backups, 
    and packet capture operations and recieves one argument:
    MuConnection - The connection class for maintaining a SSL session to 
    the Mu Test Suite

    TestCase: Not implemented yet: This class will contain a single line 
    of data to execute a test case with Studio Verify

    fault: this class is used to represent fault data that can be 
    retrieved from the API

TODO: data mine analysis archive for remediation data and add a archive class	
"""


class MuRestApi:
    """ MuRestApi: MuRestApi - This class is instantiated to run Fuzz anaysis, 
    backups, and packet capture operations and recieves one argument:
    
    Arguments: 
    MuConnection - The connection class for maintaining a SSL session 
    to the Mu Test Suite
    
    Globals:
    muConnection - The SSL connection manager
    session - containes the active session id for this instance
    VERBOSE - enabel detaield printing to screen.
    
    TODO:
    logging - Used to loog activiteis instead of VERBOSE
    archive - a class that represents a mu test result archive fro data mining
    
    """
    muConnection = None
    session = None
    VERBOSE = False

    def __init__(self, connection):
        """ set up the TestAPI with the passed in MuConnection """

        self.muConnection = connection

    def runBackup(self):
        """ initiate a backup job on the Mu Test 
        Example URL: https://$MU_IP/api/v3/backup/run
        
        Arguments:
        None
        
        Returns:
        MuConnection.MuResponseObject
        """
        
        url = self.muConnection.TOP_LEVEL_URL + '/api/v3/backup/run'
        responseData = self.muConnection.get(url)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def getBackupStatus(self, jobId):
        """ get the status a backup job on the Mu 
        Example URL: https://$MU_IP/api/v3/backup/status?jobId=1124485278
        
        Arguments:
        jobId - the jobId to get the status on
        
        Returns:
        MuConnection.MuResponseObject
        """

        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/backup/status?jobId=' + jobId
        responseData = self.muConnection.get(url)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def getBackup(self, jobId, fileName):
        """ get the completed backup job on the Mu 
        Example URL: https://$MU_IP/api/v3/backup/get?jobId=1124485278
        
        Arguments:
        jobId - the jobId to get the status on
        fileName - the string filename including path to save the backub file as
        
        Returns:
        MuConnection.MuResponseObject
        """

        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/backup/get?jobId=' + jobId
        responseData = self.muConnection.getFile(url, fileName)
        self.session = responseData
        log.debug(responseData)
        return responseData.isCompleted()

    def importTemplates(self, fileName):
        """ import a template from a local file onto the Mu 
        Example URL: https://$MU_IP/api/v3/templates/import
        
        Arguments:
        fileName - the string filename including path of the templates 
        file to post to the Mu
        
        Returns:
        MuConnection.MuResponseObject
        """

        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/templates/import'
        responseData = self.muConnection.postFile(url, fileName)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def exportTemplates(self, fileName, uuid=None, type=None, name=None):
        """ import a template from a local file onto the Mu. One must use 
        (uuid) or (type and name)
        Example URL: 
        https://$MU_IP/api/v3/templates/export?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
        Example URL: 
        https://$MU_IP/api/v3/templates/export?type=scenario&name=MyScenarioName
        
        Arguments:
        fileName - the string filename including path of the templates 
        save retrieved template from the Mu
        uuid - the uuid of the template to export
        type - the type of template to export
        name - the name of the template to export
        
        Returns:
        MuConnection.MuResponseObject
        """
        if uuid != None:
            url = self.muConnection.TOP_LEVEL_URL + \
                '/api/v3/templates/export?uuid=' + uuid
            responseData = self.muConnection.getFile(url, fileName)
            self.session = responseData
        elif type != None and name != None:
            name2 = name.replace(" ", "%20")
            url = self.muConnection.TOP_LEVEL_URL + \
                '/api/v3/templates/export?type=' + type + '&name=' + \
                name2
            responseData = self.muConnection.getFile(url, fileName)
            self.session = responseData
        else:
            log.info("invalid template get")
        log.debug(responseData)
        return responseData

    def deleteTemplates(self, uuid):
        """ delete a template from the Mu.
        Example URL: 
        https://$MU_IP/api/v3/template/delete?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
        
        Arguments:
        uuid - the uuid of the template to export
        
        Returns:
        MuConnection.MuResponseObject
        """
        
        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/template/delete?uuid=' + uuid
        responseData = self.muConnection.get(url)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def listTemplates(self, templateType):
        """ list templates on the Mu.
        Example URL: https://$MU_IP/api/v3/analysis/list
        Example URL: https://$MU_IP/api/v3/scenario/list
        
        Arguments:
        templateType - the type of mu template to list
        
        Returns:
        MuConnection.MuResponseObject
        """
        url = self.muConnection.TOP_LEVEL_URL + '/api/v3/' + \
            templateType + '/list'
        responseData = self.muConnection.get(url)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def listTemplateTypes(self):
        """ list all template types.
        Example URL: https://$MU_IP/api/v3/templates/types
        
        Arguments:
        none
        
        Returns:
        MuConnection.MuResponseObject
        """
        
        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/templates/types'
        responseData = self.muConnection.get(url)
        self.session = responseData
        if not responseData.hasError():
            tmpData = responseData.getRawResponse()
            tmpData = tmpData.split("\n")
            responseData.setRetrievedObjectType(responseData.LIST)
            responseData.setRetrievedObject(tmpData)
            responseData.setCompleted(True)
        log.debug(responseData)
        return responseData

    def runAnalysis(self, uuid):
        """ run an anlysis (DoS, PVA, Fuzz tests) on the Mu.
        Example URL: 
        https://$MU_IP/api/v3/analysis/run?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

        
        Arguments:
        uuid - the previusly run analysis uuid or the analysis template 
        uuid on the Mu
        
        Returns:
        MuConnection.MuResponseObject
        """
        
        return self.controlAnalysis("run", uuid)

    def getAnalysisStatus(self, uuid):
        """ get the status of an anlysis (DoS, PVA, Fuzz tests) on the Mu.
        Example URL: 
        https://$MU_IP/api/v3/analysis/status?uuid=7405e8af-2863-433a-9a09-d7a211a318cb

        
        Arguments:
        uuid - the analysis uuid on the Mu
        
        Returns:
        MuConnection.MuResponseObject
        """
        
        return self.controlAnalysis("status", uuid)

    def abortAnalysis(self, uuid):
        """ abort an anlysis (DoS, PVA, Fuzz tests) on the Mu.
        Example URL: 
        https://$MU_IP/api/v3/analysis/stop?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
        
        Arguments:
        uuid - the analysis uuid on the Mu
        
        Returns:
        MuConnection.MuResponseObject
        """

        return self.controlAnalysis("stop", uuid)

    def pauseAnalysis(self, uuid):
        """ pause an anlysis (DoS, PVA, Fuzz tests) on the Mu.
        Example URL: 
        https://$MU_IP/api/v3/analysis/pause?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
        
        Arguments:
        uuid - the analysis uuid on the Mu
        
        Returns:
        MuConnection.MuResponseObject
        """

        return self.controlAnalysis("pause", uuid)

    def resumeAnalysis(self, uuid):
        """ resume an anlysis (DoS, PVA, Fuzz tests) on the Mu.
        Example URL: 
        https://$MU_IP/api/v3/analysis/resume?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
        
        Arguments:
        uuid - the analysis uuid on the Mu
        
        Returns:
        MuConnection.MuResponseObject
        """

        return self.controlAnalysis("resume", uuid)

    def listAnalysis(self):
        """ list all anlysis (DoS, PVA, Fuzz tests) on the Mu.
        Example URL: https://$MU_IP/api/v3/analysis/list
        
        Arguments:
        none
        
        Returns:
        MuConnection.MuResponseObject
        """
        url = self.muConnection.TOP_LEVEL_URL + '/api/v3/analysis/list'
        responseData = self.muConnection.get(url)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def deleteAnalysis(self, uuid):
        """ delete an anlysis (DoS, PVA, Fuzz tests) on the Mu.
        Example URL: 
        https://$MU_IP/api/v3/analysis/delete?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
        
        Arguments:
        uuid - the analysis uuid on the Mu
        
        Returns:
        MuConnection.MuResponseObject
        """

        return self.controlAnalysis("delete", uuid)

    def controlAnalysis(self, action, uuid):
        """ used for all analysis operations (run,stop,pause,resume,delete) 
        on the Mu.
        Example URL: 
        https://$MU_IP/api/v3/analysis/stop?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
        Example URL: 
        https://$MU_IP/api/v3/analysis/delete?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
        
        Arguments:
        action - the control action to execute: run,stop,pause,resume,delete
        uuid - the analysis uuid on the Mu
        
        Returns:
        MuConnection.MuResponseObject
        """
        url = self.muConnection.TOP_LEVEL_URL + '/api/v3/analysis/' + \
            action + '?uuid=' + uuid
        responseData = self.muConnection.get(url)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def getAnalysisFaultList(self, uuid):
        """ gat the fault list for an anlysis (running, aborted, finished, 
        failed, paused) on the Mu.
        Example URL: 
        https://$MU_IP/api/v3/analysis/getFaultList?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
        
        Arguments:
        uuid - the analysis uuid on the Mu
        
        Returns:
        MuConnection.MuResponseObject
        """
        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/analysis/getFaultList?uuid=' + uuid
        responseData = self.muConnection.get(url)
        self.session = responseData
        if responseData.isCompleted() == False:
            return responseData
        myTree = ElementTree.XML(responseData.getRawResponse())
        for element in myTree:
            log.info(element.tag)
            if element.tag == "fault_list":
                childList = element.getchildren()
                faultList = []
                for element2 in childList:
                    if element2.tag == "tba_fault":
                        myFault = fault(element2)
                        faultList.append(myFault)
                responseData.setRetrievedObjectType(responseData.FAULT_LIST)
                responseData.setRetrievedObject(faultList)
        log.debug(responseData)
        return responseData

    def runArchive( self, uuid, title=None, notes=None, includeEVTs=True,
        includePcaps=True, includeSamples=True, includeLogo=True,
        includeAudit=True, sign=False,):
        """ start an archie job on an analysis (running, aborted, finished, 
        failed) on the Mu.
        Example URL: 
        https://$MU_IP/api/v3/archive/run?uuid=7405e8af-2863-433a-9a09-d7a211a318cb
        
        Arguments:
        uuid - the analysis uuid on the Mu
        title - the Title to place at the top of the anaylsis reports
        notes - string of data to add to notes field of reports
        includeEVTs - boolean to include or exclude evts, default = true
        includePcaps - boolean to include or exclude test result pcaps, 
        default = true
        includeSamples - boolean to include or exclude sample pcaps, 
        default = true
        includeLogo - boolean to include or exclude logo if one is loaded 
        on the Mu, default = true
        includeAudit - boolean to include or exclude audit, default = true
        sign - boolean to indicate id digital signature is eneded, 
        default  = false
        
        Returns:
        MuConnection.MuResponseObject
        """

        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/archive/run?uuid=' + uuid
        if notes != None:
            url = url + '&notes=' + notes.replace(" ", "%20")
        if title != None:
            url = url + '&title=' + title.replace(" ", "%20")
        tmpS = ''
        if includeEVTs:
            tmpS = 'evts'
        if includePcaps:
            if tmpS == '':
                tmpS = 'pcaps'
            else:
                tmpS = tmpS + ',' + 'pcaps'
        if includeSamples:
            if tmpS == '':
                tmpS = 'samples'
            else:
                tmpS = tmpS + ',' + 'samples'
        if includeLogo:
            if tmpS == '':
                tmpS = 'logo'
            else:
                tmpS = tmpS + ',' + 'logo'
        if includeAudit:
            if tmpS == '':
                tmpS = 'audit'
            else:
                tmpS = tmpS + ',' + 'audit'
        if tmpS != '':
            url = url + '&include=' + tmpS
        if sign:
            url = url + '&sign'
        responseData = self.muConnection.get(url)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def getArchiveStatus(self, jobId):
        """ get the status of an archive job on the Mu.
        Example URL: https://$MU_IP/api/v3/archive/status?jobId==7405e8af
        
        Arguments:
        jobId - the job id of the archive
        
        Returns:
        MuConnection.MuResponseObject
        """
        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/archive/status?jobId=' + jobId
        responseData = self.muConnection.get(url)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def getArchive(self, jobId, fileName):
        """ gat the completed archive job on the Mu and save to a file.
        Example URL: https://$MU_IP/api/v3/archive/get?jobId==7405e8af
        
        Arguments:
        jobId - the job id of the archive
        fileName - The filename including path to save the archive data in.
        
        Returns:
        MuConnection.MuResponseObject
        """
        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/archive/get?jobId=' + jobId
        responseData = self.muConnection.getFile(url, fileName)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def runCapture(self, port):
        """ start a packet capture job on the Mu server.
        Example URL: https://$MU_IP/api/v3/capture/run?port=b2
        
        Arguments:
        port - the port of the mu to begin capture on.
        
        Returns:
        MuConnection.MuResponseObject
        """
        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/capture/run?port=' + port
        responseData = self.muConnection.get(url)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def getCaptureStatus(self, jobId):
        """ get the capture status of a job on the Mu.
        Example URL: https://$MU_IP/api/v3/capture/status?jobId=7405e8af
        
        Arguments:
        jobId - the jobid of the capture in question
        
        Returns:
        MuConnection.MuResponseObject
        """
        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/capture/status?jobId=' + jobId
        responseData = self.muConnection.get(url)
        self.session = responseData
        log.debug(responseData)
        return responseData

    def getCapture(self, jobId, fileName):
        """ get the packets captured by the capture job on the Mu.
        Example URL: https://$MU_IP/api/v3/capture/get?jobId=7405e8af
        
        Arguments:
        jobId - the jobid of the capture in question
        fileName - The filename including path to save the packet capture 
        data in.
        
        Returns:
        MuConnection.MuResponseObject
        """
        url = self.muConnection.TOP_LEVEL_URL + \
            '/api/v3/capture/get?jobId=' + jobId
        responseData = self.muConnection.getFile(url, fileName)
        self.session = responseData
        log.debug(responseData)
        return responseData


class fault:
    """ fault:  This class is represents the data that is retreivable from 
    the xml response of the 
    getAnalysisFaults command:
    
    Arguments: 
    xmlElement - the xml data from the mu that contains all fault data 
    for an indavidual fault
    
    TODO:
    """
    def __init__(self, xmlElement=None):
        if xmlElement == None:
            return
        if xmlElement.tag == "tba_fault":
            childList = xmlElement.getchildren()
            for element2 in childList:
                if element2.tag == "confidence_level":
                    self.setConfidenceLevel(element2.text)
                elif element2.tag == "title":
                    self.setTitle(element2.text)
                elif element2.tag == "fault_detection":
                    self.setFaultDetection(element2.text)
                elif element2.tag == "fault_isolation":
                    self.setFaultIsolation(element2.text)
                elif element2.tag == "attack_type":
                    self.setAttackType(element2.text)
        else:
            return

    def __str__(self):
        s = 'Confidence Level = ' + self.confidenceLevel + '; '
        s = s + 'Title = ' + self.title + '; '
        s = s + 'Fault Isolation Method = ' + self.faultIsolation + '; '
        s = s + 'Fault Detection Method = ' + self.faultDetection + '; '
        s = s + 'Attack Type = ' + self.attackType
        return s

    ONE = '1'
    TWO = '2'
    THREE = '3'
    FOUR = '4'
    FIVE = '5'

    VECTOR = 'vector'
    VECTOR_RANGE = 'vector range'
    VECTOR_LOOP = 'vector loop'
    VARIANT_RANGE = 'variant range'

    ISTRAMENTATION = 'instramentation'
    MONITOR = 'monitor'
    CONNECTION = 'connection'

    confidenceLevel = ''
    title = ''
    faultIsolation = ''
    faultDetection = ''
    attackType = ''

    def getConfidenceLevel(self):
        return self.confidenceLevel

    def getTitle(self):
        return self.title

    def getFaultIsolation(self):
        return self.faultIsolation

    def getFaultDetection(self):
        return self.faultDetection

    def getAttackType(self):
        return self.attackType

    def setConfidenceLevel(self, obj):
        self.confidenceLevel = obj

    def setTitle(self, obj):
        self.title = obj

    def setFaultIsolation(self, obj):
        self.faultIsolation = obj

    def setFaultDetection(self, obj):
        self.faultDetection = obj

    def setAttackType(self, obj):
        self.attackType = obj


